#%%
import numpy as np
import os
import pickle
#%%
dataset_dir = os.getcwd()#Python 代码运行时的工作目录路径，dir是directory的路径
key_file = {
    'train_img':'train-images-idx3-ubyte',
    'train_label':'train-labels-idx1-ubyte',
    'test_img':'t10k-images-idx3-ubyte',
    'test_label':'t10k-labels-idx1-ubyte'
}#字典
#.pkl 是 Python 中用于存储序列化数据的文件后缀，通常与 pickle 模块配合使用，用于实现对象的持久化存储。
save_file = dataset_dir + "/mnist.pkl"#numpy类型的数据保存的名称
train_num = 60000 #训练集数目
test_num = 10000 #测试集数目
img_dim = (1, 28, 28) #表示图像的维度信息，是一个元组，表示这是一张 1 通道（灰度）、28×28 像素大小的图像
img_size = 784 #图像展平后的一维数组长度
#%%
def _load_label(file_name):
    file_path = dataset_dir + "/data/" + file_name
    print("Converting " + file_name + " to NumPy Array ...")
    with open(file_path, 'rb') as f:
            labels = np.frombuffer(f.read(), np.uint8, offset=8)
    #np.frombuffer：从缓冲区（这里是f.read()返回的二进制数据）创建数组。
    #f.read:要解析的二进制数据（f.read()的结果）。
    #np.uint8：指定数组元素的数据类型为 “无符号 8 位整数”（0~255 范围），因为标签通常是 0~9 的数字，用 1 字节存储足够。
    #第三个参数offset=8：表示跳过文件开头的 8 个字节（偏移量为 8），从第 9 个字节开始读取有效数据。这是因为 MNIST 等数据集的二进制文件有固定格式：前几个字节是 “文件头”（存储数据类型、数量等元信息），从第 8 字节后才是实际的标签数据。
    print("Done:load label")

    return labels

def _load_img(file_name):
    file_path = dataset_dir + "/data/" + file_name
    print("Converting " + file_name + " to NumPy Array ...")
    with open(file_path, 'rb') as f:
            data = np.frombuffer(f.read(), np.uint8, offset=16)#此时得到的还是二进制一维数组
    data = data.reshape(-1, img_size)
    #将data数组重塑为 “N 行、img_size 列” 的新数组，其中 N 由 NumPy 根据总元素数量自动计算。
    print("Done:load img")

    return data
#将数据集的图像和标签数据整合到一个字典中，方便后续统一处理和使用。
def _convert_numpy():
    dataset = {}#一个字典
    #字典的键是每一类数据名称，字典的值是读取到的np数组
    dataset['train_img'] =  _load_img(key_file['train_img'])
    dataset['train_label'] = _load_label(key_file['train_label'])
    dataset['test_img'] = _load_img(key_file['test_img'])
    dataset['test_label'] = _load_label(key_file['test_label'])

    return dataset
#%%
def init_mnist():
    dataset = _convert_numpy()
    print("Creating pickle file ...")
    #使用pickle.dump()函数将dataset字典序列化并写入文件：
#第一个参数dataset：要保存的 Python 对象（这里是包含数据集的字典）。
#第二个参数f：文件对象（即前面打开的save_file）。
#第三个参数-1：指定使用最高级别的压缩算法，减少.pkl文件的体积。
    with open(save_file, 'wb') as f:
        pickle.dump(dataset, f, -1)
    print("Done!")
#%%
def change_one_hot_label(label):
    T=np.zeros((label.size,10))
    for i in range(label.size):
        T[i][label[i]]=1
    return T
#%%
def load_mnist(normalize=True,flatten=True,one_hot_label=False):
    if not os.path.exists(save_file):
        init_mnist()
    with open(save_file,'rb') as f:
        dataset=pickle.load(f)
    #正则化(对图像数据）
    if normalize:
        for key in ('train_img','test_img'):
            dataset[key]=dataset[key].astype(np.float32)
            dataset[key]/=255.0
    #图热编码（对标签数据）
    if one_hot_label:
        for key in ('train_label','test_label'):
            dataset[key]=change_one_hot_label(dataset[key])
    #非展开
    if not flatten:
        for key in ('train_img','test_img'):
            dataset[key]=dataset[key].reshape(-1,1,28,28)
    return (dataset['train_img'],dataset['train_label']),(dataset['test_img'],dataset['test_label'])